/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.kyuubi.ha.client

import org.apache.zookeeper.{WatchedEvent, Watcher}

import org.apache.kyuubi.KyuubiException
import org.apache.kyuubi.service.Serverable

/**
 * A service for service discovery used by kyuubi server side.
 * We add another zk watch so that we can stop server more genteelly.
 *
 * @param name the name of the service itself
 * @param server the instance uri a service that used to publish itself
 */
class KyuubiServiceDiscovery private(
    name: String,
    server: Serverable) extends ServiceDiscovery(server) {
  def this(server: Serverable) =
    this(classOf[KyuubiServiceDiscovery].getSimpleName, server)

  override def start(): Unit = {
    super.start()
    // Set a watch on the serviceNode
    val watcher = new DeRegisterWatcher
    if (zkClient.checkExists.usingWatcher(watcher).forPath(serviceNode.getActualPath) == null) {
      // No node exists, throw exception
      throw new KyuubiException(s"Unable to create znode for this Kyuubi " +
        s"instance[${server.connectionUrl}] on ZooKeeper.")
    }
  }

  override def stopGracefully(): Unit = {
    stop()
    while (server.backendService != null &&
      server.backendService.sessionManager.getOpenSessionCount > 0) {
      Thread.sleep(1000 * 60)
    }
    server.stop()
  }

  class DeRegisterWatcher extends Watcher {
    override def process(event: WatchedEvent): Unit = {
      if (event.getType == Watcher.Event.EventType.NodeDeleted) {
        warn(s"This Kyuubi instance ${server.connectionUrl} is now de-registered from" +
          s" ZooKeeper. The server will be shut down after the last client session completes.")
        stopGracefully()
      }
    }
  }
}


